package org.fujay.commons.langchain4j.rag.processor;

import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.DocumentTransformer;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.data.segment.TextSegmentTransformer;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.Builder;
import lombok.Data;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @author slm
 * @description 文档处理器配置
 * 包含文档处理过程的各种配置参数
 */
@Data
@Builder
public class DocumentProcessorConfig {

    /**
     * 文档转换器，用于在分块前对文档进行预处理
     */
    private DocumentTransformer documentTransformer;

    /**
     * 文档分块器，用于将文档切分为适合嵌入的文本段
     */
    private DocumentSplitter documentSplitter;

    /**
     * 文本段转换器，用于在嵌入前对文本段进行处理
     */
    private TextSegmentTransformer textSegmentTransformer;

    /**
     * 嵌入模型，用于生成文本的向量表示
     */
    private EmbeddingModel embeddingModel;

    /**
     * 嵌入存储，用于存储嵌入向量和原始文本段
     */
    private EmbeddingStore<TextSegment> embeddingStore;

    /**
     * 批处理大小，控制一次批量处理的文本段数量
     */
    @Builder.Default
    private int batchSize = 16;

    /**
     * 并行处理线程数
     */
    @Builder.Default
    private int parallelism = Runtime.getRuntime().availableProcessors();

    /**
     * 线程池，用于并行处理
     */
    @Builder.Default
    private ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());

    /**
     * 是否自动关闭线程池
     */
    @Builder.Default
    private boolean autoShutdownExecutor = true;

    /**
     * 重试次数，处理失败时的最大重试次数
     */
    @Builder.Default
    private int maxRetries = 3;

    /**
     * 重试间隔（毫秒）
     */
    @Builder.Default
    private long retryIntervalMs = 1000;

    /**
     * 处理超时（毫秒），单个文档处理的最大时间
     */
    @Builder.Default
    private long timeoutMs = 60000;

    /**
     * 文档ID前缀
     */
    @Builder.Default
    private String idPrefix = "doc_";

    /**
     * 支持的文件扩展名集合
     */
    @Builder.Default
    private Set<String> supportedExtensions = getDefaultSupportedExtensions();

    /**
     * 中文分词模式
     */
    @Builder.Default
    private boolean chineseSegmentation = true;

    /**
     * 中文分词器类型
     */
    @Builder.Default
    private String chineseSegmenterType = "jieba";

    /**
     * 是否将处理后的文档元数据存储到外部数据库
     */
    @Builder.Default
    private boolean storeMetadata = false;

    /**
     * 元数据存储连接URL
     */
    private String metadataStoreUrl;

    /**
     * 命名空间
     */
    @Builder.Default
    private String nameSpace = "sys";

    /**
     * 获取默认支持的文件扩展名
     */
    private static Set<String> getDefaultSupportedExtensions() {
        Set<String> extensions = new HashSet<>();
        extensions.add("txt");
        extensions.add("md");
        extensions.add("pdf");
        extensions.add("doc");
        extensions.add("docx");
        extensions.add("ppt");
        extensions.add("pptx");
        extensions.add("xls");
        extensions.add("xlsx");
        extensions.add("html");
        extensions.add("htm");
        extensions.add("json");
        extensions.add("xml");
        extensions.add("csv");
        return extensions;
    }

    /**
     * 获取线程池（如果为null则创建新线程池）
     */
    public ExecutorService getExecutorService() {
        if (executorService == null) {
            return Executors.newFixedThreadPool(parallelism);
        }
        return executorService;
    }

}